# -*- coding: utf-8 -*-
import torch
from torch import nn
import random

"""add models"""


class tokmak_model(nn.Module):
    def __init__(
        self,
        in_channels,
        hidden_size,
        num_layers,
        dropout_rate: float = 0.0,
        out_channels: int = 1,
        noise_ratio: float = 0.1,
    ) -> None:
        """
        Args:
            noise_ratio (float): add GaussianNoise for slightly blur previous wrong encoder input (measurement error e.t.c.) and decoder input (network wrong output).
        Returns:
            None
        """
        super().__init__()
        self.RNN = nn.LSTM(
            in_channels, hidden_size, num_layers, batch_first=True, dropout=dropout_rate
        )
        self.dense = nn.Linear(hidden_size, out_channels)
        self.noise_ratio = noise_ratio

    def forward(
        self, enc_inputs, dec_inputs, state=None, teacher_forcing_ratio=0
    ):
        """
        Args:
            enc_inputs (Tensor): [batch_size, seq_len, input_size],
                encoder input of the network
            dec_inputs (Tensor): [batch_size, seq_len, output_size],
                decoder input of the network.
            teacher_forcing_ratio (Float): teaching force ratio. in our case, it will change step by step.
        Returns:
            model_ouputs (Tensor): [batch_size, seq_len, output_size]
        """
        # add noise
        if self.training:
            noise = torch.randn_like(enc_inputs) * self.noise_ratio
            enc_inputs = (
                (noise + enc_inputs).detach() - enc_inputs
            ).detach() + enc_inputs
            noise = torch.randn_like(dec_inputs) * self.noise_ratio
            dec_inputs = (
                (noise + dec_inputs).detach() - dec_inputs
            ).detach() + dec_inputs
        dec_inputs_len = dec_inputs.shape[1]
        model_outputs = []
        self.RNN.flatten_parameters()
        for t in range(0, dec_inputs_len):
            if t == 0:
                # cat_input shape [batch_size, feature_size]
                cat_input = torch.cat(
                    (enc_inputs[:, 0], dec_inputs[:, 0]), dim=-1
                )
            cat_input = cat_input.unsqueeze(1)
            # model_output shape [batch_size, 1, hidden_size]
            model_output, state = self.RNN(cat_input, state)
            model_output = self.dense(model_output)
            model_outputs.append(model_output)
            teacher_force = random.random() <= teacher_forcing_ratio
            padded_output_value = (
                dec_inputs[:, t] if teacher_force else model_output[:, 0, :]
            )
            cat_input = torch.cat((enc_inputs[:, t], padded_output_value), dim=-1)
        model_outputs = torch.cat(model_outputs, 1)
        return model_outputs, state
